from weakref import KeyedRef
import torch
from torch._C import set_flush_denormal
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.activation import PReLU
from torch.nn.modules.conv import Conv2d
from torch.nn.modules.pooling import MaxPool2d


def weights_init(m):
    if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight.data)
        nn.init.constant_(m.bias, 0.1)


class LossFn:
    def __init__(self, cls_factor=1, box_factor=1):
        # loss function
        self.cls_factor = cls_factor
        self.box_factor = box_factor
        self.loss_cls = nn.BCELoss() # binary cross entropy
        self.loss_box = nn.MSELoss() # mean square error


    def cls_loss(self,gt_label,pred_label):
        pred_label = torch.squeeze(pred_label)
        gt_label = torch.squeeze(gt_label)
        # get the mask element which >= 0, only 0 and 1 can effect the detection loss
        mask = torch.ge(gt_label,0)
        valid_gt_label = torch.masked_select(gt_label,mask)
        valid_pred_label = torch.masked_select(pred_label,mask)
        return self.loss_cls(valid_pred_label,valid_gt_label)*self.cls_factor


    def box_loss(self,gt_label,gt_offset,pred_offset):
        pred_offset = torch.squeeze(pred_offset)
        gt_offset = torch.squeeze(gt_offset)
        gt_label = torch.squeeze(gt_label)

        #get the mask element which != 0
        unmask = torch.eq(gt_label,0)
        mask = torch.eq(unmask,0)
        #convert mask to dim index
        chose_index = torch.nonzero(mask.data)
        chose_index = torch.squeeze(chose_index)
        #only valid element can effect the loss
        valid_gt_offset = gt_offset[chose_index,:]
        valid_pred_offset = pred_offset[chose_index,:]
        return self.loss_box(valid_pred_offset,valid_gt_offset)*self.box_factor

# 简化MTCNN-PNet & ONet

class PNet(nn.Module):
    '''
    input:
    12x47x3

    output:
    classifcation:1x1x2
    bbox regression:1x1x4
    '''

    def __init__(self) -> None:
        super(PNet, self).__init__()

        # backend
        self.pre_layer = nn.Sequential(
            nn.Conv2d(3, 10, kernel_size=3, stride=1),  # conv1
            nn.PReLU(),  # PReLU1
            nn.MaxPool2d(kernel_size=(2, 5), stride=(2, 5)),    # pool1
            nn.Conv2d(10, 16, kernel_size=(3, 5), stride=(1, 1)),   # conv2
            nn.PReLU(),  # PReLU2
            nn.Conv2d(16, 32, kernel_size=(3, 5), stride=1),  # conv3
            nn.PReLU()  # PReLU3
        )
        # detection / classification
        self.conv4_1 = nn.Conv2d(32, 2, kernel_size=1, stride=1)
        # bbox regression
        self.conv4_2 = nn.Conv2d(32, 4, kernel_size=1, stride=1)
        # weight initiation with xavier
        self.apply(weights_init)

    def forward(self, x):
        x = self.pre_layer(x)
        # detection
        label = F.sigmoid(self.conv4_1(x))
        # regression
        offset = self.conv4_2(x)

        return label, offset


class ONet(nn.Module):
    '''
    input:
    24x94x3

    output:
    classifcation:1x1x2
    bbox regression:1x1x4
    '''

    def __init__(self) -> None:
        super(ONet, self).__init__()

        # backend
        self.pre_layer = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1),   # conv1
            nn.PReLU(),  # PRuLU1
            nn.MaxPool2d(kernel_size=2, stride=2),  # pool1
            nn.Conv2d(32, 64, kernel_size=3, stride=1),  # conv2
            nn.PReLU(),  # PReLU2
            nn.MaxPool2d(kernel_size=2, stride=2),  # pool2
            nn.Conv2d(64, 64, kernel_size=3, stride=1),  # conv3
            nn.PReLU(),  # PReLU3
            nn.MaxPool2d(kernel_size=2, stride=2),  # pool3
            nn.Conv2d(64, 128, kernel_size=1, stride=1),    # conv4
            nn.PReLU()  # PReLU4
        )
        self.fc1 = nn.Linear(128*10*1, 256)  # dense1
        self.PReLU5 = nn.PReLU()

        # detection / classification
        self.fc2_1 = nn.Linear(256, 2)
        # bbox regression
        self.fc2_2 = nn.Linear(256, 4)
        # weight initiation with xavier
        self.apply(weights_init)

    def forward(self, x):
        # backend
        x = self.pre_layer(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.PReLU5(x)
        # detection
        det = F.sigmoid(self.fc2_1(x))
        # bbox regression
        bbox = self.fc2_2(x)

        return det, bbox

if __name__ == '__main__':
    P = PNet()
    for para in P.parameters():
        print(para)